#include "protocol_api.h"
#include "protocol_board.h"
#include "protocol_encryption.h"

uint16_t ProtocolEncryption(void *pOriginalData)
{
  uint8_t *pDataBuffOffset = NULL;
  uint8_t KeyValue[ENCRYPTION_KEY_LENGTH] = {0};
  uint16_t EncryptedDataLength = 0, MessageSize = 0;
  ProtocolFormat_t *pProtocolFormat = (ProtocolFormat_t *)pOriginalData;

  if (pOriginalData != NULL)
  {
    ProtocolGetEncryptionKey(KeyValue);
    EncryptedDataLength = ((pProtocolFormat->Length[0] << 8) | pProtocolFormat->Length[1]) + PROTOCOL_CAPTCHA_BIT_LENGTH + PROTOCOL_VERSION_BIT_LENGTH + PROTOCOL_CMD_BIT_LENGTH;
    if ((EncryptedDataLength % 16) != 0)
    {
      EncryptedDataLength += (16 - (EncryptedDataLength % 16));
    }

    pDataBuffOffset = pProtocolFormat->Captcha;
    for (uint8_t i = 0; i < EncryptedDataLength / 16; i++)
    {
      ProtocolCalculateAesEcbEncrypt(KeyValue, pDataBuffOffset);
      pDataBuffOffset += 16;
    }

    ProtocolGetOuterLayerEncryptionKey(KeyValue);
    EncryptedDataLength += PROTOCOL_RANDOM_BIT_LENGTH + PROTOCOL_TIMESTAMP_BIT_LENGTH + PROTOCOL_INDEX_BIT_LENGTH;
    if ((EncryptedDataLength % 16) != 0)
    {
      EncryptedDataLength += (16 - (EncryptedDataLength % 16));
    }

    pDataBuffOffset = pProtocolFormat->Random;
    for (uint8_t i = 0; i < EncryptedDataLength / 16; i++)
    {
      ProtocolCalculateAesEcbEncrypt(KeyValue, pDataBuffOffset);
      pDataBuffOffset += 16;
    }

    MessageSize = (EncryptedDataLength - (PROTOCOL_FIXED_BIT_LENGTH - (PROTOCOL_SYNCWORD_BIT_LENGTH + PROTOCOL_DATA_SIZE_BIT_LENGTH)));
    pProtocolFormat->Length[0] = (uint8_t)((MessageSize & 0xFF00) >> 8);
    pProtocolFormat->Length[1] = (uint8_t)(MessageSize & 0x00FF);
  }

  return EncryptedDataLength;
}

void ProtocolDecryption(void *pOriginalData)
{
  if (pOriginalData != NULL)
  {
    uint16_t DecryptDataLength = 0;
    uint8_t *pDataBuffOffset = NULL;
    uint8_t KeyValue[ENCRYPTION_KEY_LENGTH] = {0};
    ProtocolFormat_t *pProtocolFormat = (ProtocolFormat_t *)pOriginalData;

    ProtocolGetOuterLayerEncryptionKey(KeyValue);
    DecryptDataLength = ((pProtocolFormat->Length[0] << 8) | pProtocolFormat->Length[1]) + (PROTOCOL_FIXED_BIT_LENGTH - (PROTOCOL_SYNCWORD_BIT_LENGTH + PROTOCOL_DATA_SIZE_BIT_LENGTH));
    if ((DecryptDataLength % 16) == 0)
    {
      pDataBuffOffset = pProtocolFormat->Random;
      for (uint8_t i = 0; i < DecryptDataLength / 16; i++)
      {
        ProtocolCalculateAesEcbDecrypt(KeyValue, pDataBuffOffset);
        pDataBuffOffset += 16;
      }

      ProtocolGetEncryptionKey(KeyValue);
      DecryptDataLength -= (PROTOCOL_RANDOM_BIT_LENGTH + PROTOCOL_TIMESTAMP_BIT_LENGTH + PROTOCOL_INDEX_BIT_LENGTH);
      if ((DecryptDataLength % 16) == 0)
      {
        pDataBuffOffset = pProtocolFormat->Captcha;
        for (uint8_t i = 0; i < DecryptDataLength / 16; i++)
        {
          ProtocolCalculateAesEcbDecrypt(KeyValue, pDataBuffOffset);
          pDataBuffOffset += 16;
        }
      }
    }
  }
}
